#coding=utf-8
'''
简单的光线追踪，只考虑反射
'''
import matplotlib.pyplot as plt
import pyopencl as cl
import numpy as np
import time
import os
from os import path
os.environ['PYOPENCL_COMPILER_OUTPUT'] = '0'
# Randomly generate a number between [0, x)
rnd = lambda x: x*np.random.rand()
# 实体的数量
LIGHTS=1
INF=2e10
#7680*1159
GX=512
GY=512

def f4(a):
 return np.array([a[0],a[1],a[2],0],dtype=np.float32)
def rng():
 return np.random.uniform(0,1,1)[0]
'''读取模型数据'''

def flaten(node,coun,bvh):
 coun[0]+=1
 current=coun[0]
 if node.inside: #内部节点，有两个子树
  flaten(node.left,coun,bvh) #左
  n1=flaten(node.right,coun,bvh) #右
  box1,box2=node.box #.box.reshape((6,)).tolist()
  tmp=np.zeros(12,dtype=np.float32)  #为了数据对齐
  tmp[0:3]=box1[:]
  #tmp[3]=n1
  tmp[4:7]=box2[:]
  tmp[8:12]=[n1,0,node.axis,node.ptype]
  bvh[current]=tmp
 else:
  box1,box2=node.box #.box.reshape((6,)).tolist()
  tmp=np.zeros(12,dtype=np.float32)
  tmp[0:3]=box1[:]
  tmp[4:7]=box2[:]
  tmp[8:12]=[node.offset,node.nPrimitives,node.axis,node.ptype]
  bvh[current]=tmp
 return current

def loadmodel(obj):
 # 读取材质
 index=0
 mtl={}
 try:
  fm=open(obj.replace('obj','mtl'),'r',encoding='UTF-8')
  ls=fm.readlines() 
  fm.close()
  name=''
  for line in ls:
   if line[:6]=='newmtl':
    index+=1
    name=line.split()[1]
    mtl[name]={}
    mtl[name]['index']=index-1
    continue
   if index>0:
    item=line.split()
    if len(item)>1:
     mtl[name][item[0]]=[float(i) for i in item[1:]]
  print(mtl)
 except:
  mtl['00__default']={}
  mtl['00__default']['index']=0  
  print('no mtl file, use default')
 # 读取几何
 f=open(obj,'r',encoding='UTF-8')
 ls=f.readlines()
 #print(len(ls))
 f.close()
 points=[]
 normals=[]
 faces=[]
 mtl_name='00__default'
 mtls=[]
 PrimativeType=[]
 for line in ls:
  if line[:6]=='usemtl':
   mtl_name=line.split()[1]
  if line[:2]=='v ':
   v=[float(i) for i in line[2:].split()]+[0] #凑够float4
   points.append(v)
  elif line[:2]=='vn':
   v=[float(i) for i in line[2:].split()]+[0] #凑够float4
   normals.append(v)
  #三角形面片
  elif line[:2]=='f ':
   PrimativeType.append(0)
   mtls.append(mtl_name)
   '''f 13845/13845/11017 1799/1799/11018 17864/17864/11019 '''
   p3=[i.split('/') for i in line[2:].split()]
   p=[points[int(p3[i][0])-1] for i in range(3)]
   n=[normals[int(p3[i][2])-1] for i in range(3)]
   faces.append(p+n)
  #球 c x y z r 作为一个特殊的表面
  elif line[:2]=='c ':
   PrimativeType.append(1)
   mtls.append(mtl_name)
   try:
    circle=[[float(i) for i in line[2:].split()]]*6 # x y z r
   except:
    print(line,line[2:].split())
    a=input()
   faces.append(circle)
 #print(mtls)
 print('points=',len(points))
 print('normals=',len(normals))
 print('faces=',len(faces))
 faces=np.array(faces)
 info=[[] for i in range(len(faces))]
 for i in range(len(faces)):
  #face=faces[i]
  if PrimativeType[i]==0:
   _min=np.array([np.min(faces[i][0:3,0]),np.min(faces[i][0:3,1]),np.min(faces[i][0:3,2])])
   _max=np.array([np.max(faces[i][0:3,0]),np.max(faces[i][0:3,1]),np.max(faces[i][0:3,2])])
  elif PrimativeType[i]==1:
   _min=faces[i][0,0:3]-faces[i][0,3]
   _max=faces[i][0,0:3]+faces[i][0,3]
   print(_max,_min)
  bound=[i,np.array([_min,_max]),mtls[i],(_min+_max)*0.5,PrimativeType[i]]
  info[i]=bound    
 start=0
 end=len(info)
 oderinfo=[]
 counter=[0]
 n=NODE(info,start,end,oderinfo,counter,faces,mtl)
 # 重排oderinfo: v n v n v n
 for i in range(len(oderinfo)):
  tmp=oderinfo[i]
  oderinfo[i]=[tmp[0],tmp[3],tmp[1],tmp[4],tmp[2],tmp[5]]
 oderinfo=np.array(oderinfo,dtype=np.float32)
 # 材质
 mMtl=np.zeros((3*len(mtl),4),dtype=np.float32)
 for i in range(len(mtl)-1,-1,-1):
  m=mtl.popitem()[1]
  assert i==m['index']
  mMtl[i*3,0:3]=m['diff']
  mMtl[i*3,3:4]=m['rough']
  mMtl[i*3+1,0:3]=m['spec']
  mMtl[i*3+1,3:4]=m['ior']
  mMtl[i*3+2,0:3]=m['emit'] 
 return len(oderinfo),n.getdata(counter[0]),oderinfo,mMtl

def union(a,b):
 a1=np.array([a[0],b[0]])
 a2=np.array([a[1],b[1]])
 c=np.array([[np.min(a1[:,0]),np.min(a1[:,1]),np.min(a1[:,2])],
                   [np.max(a2[:,0]),np.max(a2[:,1]),np.max(a2[:,2])]])
 return c

def surface(a):
 b=a[1]-a[0]
 return 2*(b[0]*b[1]+b[0]*b[2]+b[1]*b[2])

class NODE(object):
 left=[]
 right=[]
 box=[]
 inside=True
 nPrimitives=0
 offset=0
 axis=0
 ptype=0
 def __init__(self,info,start,end,oderinfo,counter,faces,mtls):
  counter[0]+=1
  nPrimitives = end - start
  #print(end , start,end - start)
  offset=len(oderinfo)
  if nPrimitives==1:
   self.inside=False
   tmp=info[start][1]
   self.nPrimitives=nPrimitives
   for i in range(start,end):
    num,bound,mtlname,cent,ptype=info[i]
    self.offset=offset
    self.ptype=ptype
    face=faces[num]
    face[0][3]=mtls[mtlname]['type'][0] #材料类型
    face[1][3]=mtls[mtlname]['index'] #材料索引
    oderinfo.append(face)
    try:
     tmp=union(tmp,bound)
    except:
     print(bound)
     aa=input('error')
   self.box=tmp
  else:
   #排序,根据最大延展方向
   inf=np.inf
   bound=np.array([[inf,inf,inf],[-inf,-inf,-inf]])
   boundcent=np.array([[inf,inf,inf],[-inf,-inf,-inf]])
   for i in range(start,end):
    bound=union(bound,info[i][1])
    boundcent=union(boundcent,[info[i][3],info[i][3]])
   #找到最大方向
   dim=(boundcent[1]-boundcent[0]).argmax()
   #几个个面片的包围盒完全重合(例如一个矩形分成两个三角形)
   if boundcent[1][dim]==boundcent[0][dim]:
    self.nPrimitives=nPrimitives
    self.offset=offset
    self.inside=False
    self.box=bound
    for i in range(start,end):
     num,bound,mtlname,cent,ptype=info[i]
     self.ptype=ptype
     self.offset=offset
     face=faces[num]
     face[0][3]=mtls[mtlname]['type'][0] #材料类型
     face[1][3]=mtls[mtlname]['index'] #材料索引
     oderinfo.append(face)
   else:
    #找出分割点middle
    middle=0
     #排序
    info[start:end]=sorted(info[start:end],key=lambda s:s[3][dim])
    mi=0
    if nPrimitives<=2:
     middle=start+1
    else:
     nBuckets = 12
     bt=np.array([[inf,inf,inf],[-inf,-inf,-inf]])
     buckets=[[0,bt] for i in range(nBuckets)] #空间离散化成12分
     # 计算每个分空间的bound和实体数量
     num=end-start
     
     dim_l=(boundcent[1][dim]-boundcent[0][dim])
     
     for i in range(start,end):
      n=int(nBuckets*(-boundcent[0][dim]+info[i][3][dim])/dim_l)
      n=n if n<nBuckets else nBuckets-1
      assert n>=0 and n<nBuckets
      buckets[n][0]+=1 #桶的实体加一
      buckets[n][1]=union(buckets[n][1], info[i][1]) #桶的包围盒
     
     # 计算最小代价分割 注意n个桶有n-1个划分法
     cost=[0 for i in range(nBuckets - 1)]
     for i in range(0,nBuckets-1):
      b1=np.array([[inf,inf,inf],[-inf,-inf,-inf]])
      c1=0
      for j in range(0,i+1):
       c1+=buckets[j][0]
       b1=union(b1,buckets[j][1])
      b2=np.array([[inf,inf,inf],[-inf,-inf,-inf]])
      c2=0
      for j in range(i+1,nBuckets):
       c2+=buckets[j][0]
       b2=union(b2,buckets[j][1])
      cost[i]=1.0+(c1*surface(b1)+c2*surface(b2))/surface(bound)
     
     # 找出最小划分
     #print(buckets)
     #print(cost)
     mi=cost[0]
     mid=1
     for i in range(1,nBuckets-1):
      if mi>cost[i]:
       mi=cost[i]
       mid=i
     max_dim=boundcent[0][dim]+dim_l*mid/nBuckets
     # 找出分割点
     for i in range(start,end):
      if info[i][3][dim]>max_dim:
       middle=i
       break
     
     #print('her',start,end,middle)
    if nPrimitives<=4 and mi>nPrimitives:
     self.nPrimitives=nPrimitives
     self.offset=offset
     self.inside=False
     self.box=bound     
     for i in range(start,end):
      num,bound,mtlname,cent,ptype=info[i]
      self.ptype=ptype
      self.offset=offset
      face=faces[num]
      face[0][3]=mtls[mtlname]['type'][0] #材料类型
      face[1][3]=mtls[mtlname]['index'] #材料索引
      face[2][3]=ptype #是三角形还是圆
      oderinfo.append(face)
    # # 按最小划分建立左右两个子树
    else:
     self.axis=dim
     self.left=NODE(info,start,middle,oderinfo,counter,faces,mtls)
     self.right=NODE(info,middle,end,oderinfo,counter,faces,mtls)
  self.box=self.getbox()
   
 def getbox(self):
  if self.inside==False:
   return self.box
  return union(self.left.box,self.right.box)
  
 def getdata(self,num):
  coun=[-1]
  bvh=[[] for i in range(num)]   
  flaten(self,coun,bvh)
  return np.array(bvh,dtype=np.float32)
 
def sizeofbite(arr):
  x,y=arr.shape
  return x*y*32/1024/1024
class RayTracing(object):
 def get_data(self,obj):
  name=obj.split('.')[0]+'_faces_axis.gz'
  name2=obj.split('.')[0]+'_bvh_axis.gz'
  name3=obj.split('.')[0]+'_mtl_axis.gz'
  if path.isfile(name) and path.isfile(name2) and path.isfile(name3):
   print('load file...')
   self.faces=np.loadtxt(name,dtype=np.float32)
   self.bvh=np.loadtxt(name2,dtype=np.float32)
   self.mMtl=np.loadtxt(name3,dtype=np.float32)
   self.ENTITYS=len(self.faces)
   print('内存占用= %.2f M'%(sizeofbite(self.faces)+sizeofbite(self.bvh)))
  else:
   print('build file...')
   self.ENTITYS,self.bvh,self.faces,self.mMtl=loadmodel(obj)
   name=obj.split('.')[0]
   np.savetxt(name+'_bvh_axis.gz',self.bvh)
   np.savetxt(name+'_faces_axis.gz',self.faces.reshape((len(self.faces),24)))
   np.savetxt(name+'_mtl_axis.gz',self.mMtl)
  print(self.mMtl)
 def init(self):
  '''初始化opencl环境'''
  platforms = cl.get_platforms() #获取可用的opencl平台
  print("%d 个可用的opencl平台"%(len(platforms)))
  for pf in platforms:
   print(pf.get_info(cl.platform_info.NAME),pf.get_info(cl.platform_info.VERSION))
  self.ctx = cl.Context(dev_type=cl.device_type.GPU,properties=[(cl.context_properties.PLATFORM, platforms[1])])
  print("当前使用平台：",platforms[1].get_info(cl.platform_info.NAME))
  self.queue=cl.CommandQueue(self.ctx) #opencl命令队列
  
  # #####################  场景 #####################################
  
  #self.ENTITYS=self.testscene2()
  
  #self.ENTITYS=self.randomscene_with_plate2()
  #NORMAL_UP_X,NORMAL_UP_Y,NORMAL_UP_Z=0,0,1
  #LOOK_AT_X,LOOK_AT_Y,LOOK_AT_Z=0,23,-20
  #LOOK_TO_X,LOOK_TO_Y,LOOK_TO_Z=0,0,-20
  #self.ENTITYS=1 #,bvh,faces=loadmodel()
  NORMAL_UP_X,NORMAL_UP_Y,NORMAL_UP_Z=0,1,0
  LOOK_AT_X,LOOK_AT_Y,LOOK_AT_Z=278, 278, -800
  LOOK_TO_X,LOOK_TO_Y,LOOK_TO_Z=278, 278, 0
  # ##################### 定义灯光 ###########位置 颜色 #############
  self.mLights=np.array([5., 5., -10.,0,1.0,1.0,0.0,0.0],dtype=np.float32) 
  # ##################### 定义材质 ##################################
  
  #编译内核程序
  opt='-cl-fast-relaxed-math' #'-cl-std=CL2.0' # 
  #f=open(path.dirname('__file__')+'\\'+'ray_漫反射_镜面反射_折射.c',encoding='UTF-8')
  f=open('ray_漫反射_镜面反射_折射_CornellBox.c',encoding='UTF-8')
  src=f.read() #
  f.close()
  # 替换关键字
  src=src.replace('ENTITYS_NUM',str(self.ENTITYS))
  src=src.replace('DIM_X',str(GX))
  src=src.replace('DIM_Y',str(GY))
  src=src.replace('NORMAL_UP_X',str(NORMAL_UP_X))
  src=src.replace('NORMAL_UP_Y',str(NORMAL_UP_Y))
  src=src.replace('NORMAL_UP_Z',str(NORMAL_UP_Z))
  src=src.replace('LOOK_AT_X',str(LOOK_AT_X))
  src=src.replace('LOOK_AT_Y',str(LOOK_AT_Y))
  src=src.replace('LOOK_AT_Z',str(LOOK_AT_Z))
  src=src.replace('LOOK_TO_X',str(LOOK_TO_X))
  src=src.replace('LOOK_TO_Y',str(LOOK_TO_Y))
  src=src.replace('LOOK_TO_Z',str(LOOK_TO_Z))  
  self.prg=cl.Program(self.ctx, src).build(opt) #src:内核源代码
  self.ray_tracing=self.prg.ray_tracing
 
  self.mScene	=np.zeros([GY,GX,4],dtype=np.uint8)
  mf = cl.mem_flags
  ctx=self.ctx
  print('bvh 节点数=',len(self.bvh))
  self.d_mBVH= cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=self.bvh)  
  self.d_mFaces= cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=self.faces)  
  self.d_mLights= cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=self.mLights)
  self.d_mMtl=cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=self.mMtl)  
  #  #image
  imf=cl.ImageFormat(cl.channel_order.RGBA,cl.channel_type.UNSIGNED_INT8)
  self.d_mScene=cl.Image(ctx, cl.mem_flags.WRITE_ONLY, imf, shape=(GX,GY))  
  #self.d_mScene=cl.Buffer(ctx, mf.WRITE_ONLY | mf.COPY_HOST_PTR, hostbuf=self.mScene)
  
  self.debugarr=np.zeros(64,dtype=np.float32)
  self.d_debugarr=cl.Buffer(ctx, mf.WRITE_ONLY | mf.COPY_HOST_PTR, hostbuf=self.debugarr)
  #cl.enqueue_nd_range_kernel(self.queue, self.k_initRNG,(config.PARTICLE_NUM,),(64,))
  self.ray_tracing.set_args(self.d_mFaces,self.d_mBVH,self.d_mMtl,self.d_mLights,self.d_mScene,self.d_debugarr)
 
 def run(self):
  cl.enqueue_nd_range_kernel(self.queue, self.ray_tracing,(GX,GY,),(8,8,))
  #cl.enqueue_copy(self.queue, self.mScene, self.d_mScene)
  cl.enqueue_copy(self.queue,self.mScene,self.d_mScene,origin=(0,0),region=(GX,GY))
 def show(self):
  fig, axs = plt.subplots(1, 1)
  #self.debugarr+=0.99999
  
   #bitmap = np.transpose(bitmap/255., (1, 0, 2)) # swap image's x-y axes
  
  axs.imshow(self.mScene)
  axs.grid(False)
 
   # sort the spheres by Z for visualizing z-level order, plot using circle artists
  #for ax in axs:
  axs.set_xlim(0, GX)
  axs.set_ylim(0, GY)
  #print(self.debugarr) 
  #np.savetxt('c:/lines.txt',self.debugarr)
  plt.imsave('fig2.png', np.rot90(self.mScene,k=2))
  plt.tight_layout()
  plt.show()  
 def debug(self):
  cl.enqueue_copy(self.queue, self.debugarr, self.d_debugarr)
 def savetovtk(self):
  cl.enqueue_copy(self.queue, self.debugarr, self.d_debugarr)
  vtk=open("LINES.vtk",'w')
  print("# vtk DataFile Version 1.0",file=vtk)
  print("Line representation of vtk",file=vtk)
  print("ASCII",file=vtk)
  print("",file=vtk)
  print("DATASET POLYDATA",file=vtk)
  arr=self.debugarr.reshape((1024,4))
  num=min(512,int(arr[0,3]))
  print("POINTS %d float"%(2*num),file=vtk)
  for i in range(2*num):
   print("%f %f %f"%(arr[i,0],arr[i,1],arr[i,2]),file=vtk)
  print('',file=vtk)
  print("LINES %d %d"%(2*num,num*3),file=vtk)
  for i in range(num):
   print('2 %d %d'%(i*2,i*2+1),file=vtk)
  vtk.close()
 
if __name__ == '__main__':

 p=RayTracing()
 
 #p.randomscene()
 p.get_data('CornellBox.obj')
 p.init()
 t0=time.time()
 p.run()
 print('time clapsed=',time.time()-t0)
 #p.savetovtk()
 p.show()